/**
 * @file   Tree.cpp
 * @author xingyifan
 * @date   2021/08/19 14:37:54
 * 
 * @brief  functions
 * 
 * 
 */

#include "RBT.h"

//初始化静态本地变量
Rnode *RBT::nil = NULL;
u_int RBT::count = 0;

/**
 * @file   Tree.cpp
 * @author xingyifan
 * @date   2021/08/19 22:17:08
 * 
 * @brief  RBT tree's functions
 * 
 * 
 */

Rnode *RBT::get(int depth, int i) const
{
    //i从0开始，depth表示第几层
    Rnode *p = root;
    //获取全局位置
    double left = 0;
    double right = 1 << (depth - 1);
    int floor = 1;
    while (floor != depth && p != nil)
    {
        if (i < (left + right) / 2)
        {
            p = p->lc;
            right = (left + right) / 2;
        }
        else
        {
            p = p->rc;
            left = (left + right) / 2;
        }
        floor++;
    }
    return p;
}

int RBT::indexOf(T value)
{
    if (root == nil)
    {
        return NONE;
    }
    Rnode *p = root;
    int index = p->lc->size;
    while (p != nil)
    {
        if (p->value < value)
        {
            p = p->rc;
            index += p->lc->size + 1;
        }
        else if (p->value > value)
        {
            p = p->lc;
            index -= p->rc->size + 1;
        }
        else
        {
            return index;
        }
    }
    return NONE;
}

void RBT::insert(T value)
{
    size++;
    if (root == nil)
    {
        root = new Rnode;
        root->value = value;
        root->parent = nil;
        root->lc = nil;
        root->rc = nil;
        root->size = 1;
        root->color = Black;
        return;
    }
    Rnode *p = root;
    Rnode *pos = nil; //插入位置
    while (true)
    {
        p->size++;
        if (p->value <= value)
        {
            if (p->rc != nil)
            {
                p = p->rc;
            }
            else
            {
                p->rc = new Rnode;
                pos = p->rc;
                break;
            }
        }
        else
        {
            if (p->lc != nil)
            {
                p = p->lc;
            }
            else
            {
                p->lc = new Rnode;
                pos = p->lc;
                break;
            }
        }
    }
    //赋值
    pos->value = value;
    pos->parent = p;
    pos->rc = nil;
    pos->lc = nil;
    pos->size = 1;
    //插入一个红节点，不改变黑高
    pos->color = Red;
    //调整插入节点
    adjust_insert(pos);
}

void RBT::remove(T value)
{
    Rnode *p = find_node(value);
    if (p == nil)
    {
        return;
    }
    Rnode *lp = p->lc;
    Rnode *rp = p->rc;
    Rnode *adj = nil;
    Color origin_color = p->color;
    //size调整
    change_size(root, p);
    if (rp == nil)
    {
        adj = lp;
        //用lp替换p
        replace(p, lp);
        //p大小不变，要减去1
        lp->size = p->size - 1;
        delete p;
    }
    else if (lp == nil)
    {
        adj = rp;
        //用rp替换p
        replace(p, rp);
        //p大小不变，要减去1
        rp->size = p->size - 1;
        delete p;
    }
    else
    {
        //if has rp, choose the min of rp as the new node
        Rnode *pos = min(rp);
        origin_color = pos->color;
        adj = pos->rc;
        //size调整
        change_size(p, pos);

        //if new node is not rp, link the new's parent and new's rc
        if (pos != rp)
        {
            replace(pos, pos->rc);
            pos->rc = rp;
            rp->parent = pos;
        }
        //link the new and p's parent and p's lc, rc
        pos->lc = lp;
        lp->parent = pos;

        replace(p, pos);
        pos->size = p->size;
        pos->color = p->color;
        delete p;
    }
    if (origin_color == Black)
    {
        adjust_remove(adj);
    }
}

void RBT::draw() const
{
    int h = get_h();
    //打印第一行
    int blank = (1 << h) - 2;
    print_blank(blank);
    cout << root->value << endl;

    //每行前面的空格数 2^n - 2 个
    //从第二行开始
    for (int i = 1; i < h; i++)
    {
        const int num = 1 << i;
        int c = num;
        Rnode **record = new Rnode *[num];
        blank = (1 << (h - i)) - 2;
        //前置空格
        print_blank(blank);

        //输出树枝行
        while (c > 0)
        {
            record[num - c] = get(i + 1, num - c);
            blank = (1 << (h - i)) - 1;
            //提前记录此行元素
            Rnode *p = record[num - c];
            if (p != nil)
            {
                //左端点
                if (c % 2 == 0)
                {
                    cout << "|";
                    print_(blank + 1);
                }
                else
                {
                    //回退一格，覆盖左端点为nil时产生的空格
                    cout << "\b-";
                    print_(blank);
                    cout << "|";
                    print_blank(2 * blank + 1);
                }
            }
            else if (c > 1)
            {
                //左端点
                if (c % 2 == 0)
                {
                    print_blank(blank + 2);
                }
                else
                {
                    print_blank(3 * blank + 2);
                }
            }
            c--;
        }
        cout << endl;

        c = num;
        blank = (1 << (h - i)) - 2;
        int d = 0;
        //输出元素行
        while (c > 0)
        {
            print_blank(blank);
            print_b(d);
            /*
            * 只需要在这里插入节点输出
            *
            */
            if (record[num - c] != nil)
            {
                //将可能的浮点型直接转为整型，便于格式化
                int int_value = record[num - c]->value;
                Color color = record[num - c]->color;
                cout << int_value;
                if (color == Black)
                {
                    cout << "b";
                }
                else
                {
                    cout << "r";
                }
                d = digit(int_value);
            }
            else
            {
                cout << " ";
                d = 0;
            }
            //相邻间距 2^(n+1) - 1
            blank = (1 << (h - i + 1)) - 1;
            c--;
        }
        delete[] record;
        cout << endl;
    }
}

/**
 * @file   Tree.cpp
 * @author xingyifan
 * @date   2021/08/19 15:35:42
 * 
 * @brief  private functions
 * 
 * 
 */

void RBT::inorder_walk(ostream &out, Rnode *node) const
{
    if (node != nil)
    {
        inorder_walk(out, node->lc);
        out << node->value << " ";
        inorder_walk(out, node->rc);
    }
}

Rnode *RBT::successor(Rnode *p)
{
    //if has right, return the min of right
    if (p->rc != nil)
    {
        return min(p->rc);
    }
    //if no right, find parent until pos==nil or pos comes from left
    Rnode *pos = p->parent;
    while (pos != nil && pos->lc != p)
    {
        p = pos;
        pos = p->parent;
    }
    return pos;
}

Rnode *RBT::predeccessor(Rnode *p)
{
    //if has left, return the max of left
    if (p->lc != nil)
    {
        return max(p->lc);
    }
    //if no left, find parent until pos==nil or pos comes from right
    Rnode *pos = p->parent;
    while (pos != nil && pos->rc != p)
    {
        p = pos;
        pos = p->parent;
    }
    return pos;
}

Rnode *RBT::max(Rnode *p)
{
    if (p == nil)
    {
        return nil;
    }
    while (p->rc != nil)
    {
        p = p->rc;
    }
    return p;
}

Rnode *RBT::min(Rnode *p)
{
    if (p == nil)
    {
        return nil;
    }
    while (p->lc != nil)
    {
        p = p->lc;
    }
    return p;
}

Rnode *RBT::find_node(T value)
{
    Rnode *p = root;
    while (p != nil)
    {
        if (p->value < value)
        {
            p = p->rc;
        }
        else if (p->value > value)
        {
            p = p->lc;
        }
        else
        {
            break;
        }
    }
    return p;
}

void RBT::change_size(Rnode *head, Rnode *tail)
{
    while (tail != head)
    {
        tail = tail->parent;
        tail->size--;
    }
}

void RBT::replace(Rnode *u, Rnode *v)
{
    if (u == root)
    {
        root = v;
    }
    else if (u == u->parent->lc)
    {
        u->parent->lc = v;
    }
    else
    {
        u->parent->rc = v;
    }
    v->parent = u->parent;
}

int RBT::height(Rnode *node) const
{
    if (node != nil)
    {
        return MAX(height(node->lc), height(node->rc)) + 1;
    }
    return 0;
}

void RBT::clean(Rnode *node)
{
    if (node != nil)
    {
        clean(node->lc);
        clean(node->rc);
        delete node;
        node = nil;
    }
}

int RBT::bh(Rnode *node) const
{
    if (node != nil)
    {
        int h = MAX(bh(node->lc), bh(node->rc));
        if (node->color == Black)
        {
            h++;
        }
        return h;
    }
    return 0;
}

void RBT::adjust_insert(Rnode *node)
{
    Rnode *me = node;
    Rnode *father = me->parent;
    //when parent is Red, loop
    while (father->color == Red)
    {
        //父节点是左节点
        if (father->parent->lc == father)
        {
            Rnode *uncle = father->parent->rc;
            //case 1:parent and uncle are red
            //说明父节点的父节点是黑节点
            if (uncle->color == Red)
            {
                father->color = Black;
                uncle->color = Black;
                me = father->parent;
                father = me->parent;
                me->color = Red;
            }
            else
            {
                //case 2:uncle is black
                if (father->lc != me)
                {
                    left_rotate(father);
                    me = me->lc;
                    father = me->parent;
                }
                father->color = Black;
                father->parent->color = Red;
                right_rotate(father->parent);
            }
        }
        else
        {
            Rnode *uncle = father->parent->lc;
            //case 1:parent and uncle are red
            if (uncle->color == Red)
            {
                father->color = Black;
                uncle->color = Black;
                me = father->parent;
                father = me->parent;
                me->color = Red;
            }
            else
            {
                //case 2:uncle is black
                if (father->rc != me)
                {
                    right_rotate(father);
                    me = me->rc;
                    father = me->parent;
                }
                father->color = Black;
                father->parent->color = Red;
                left_rotate(father->parent);
            }
        }
    }
    root->color = Black;
}

void RBT::adjust_remove(Rnode *node)
{
    //if me is red, change it to black, else loop
    //if me is root, the same
    Rnode *me = node;
    while (me != root && me->color == Black)
    {
        //me is the left child
        if (me == me->parent->lc)
        {
            Rnode *brother = me->parent->rc;
            //if brother is red, change it to black
            if (brother->color == Red)
            {
                me->parent->color = Red;
                brother->color = Black;
                left_rotate(me->parent);
                brother = me->parent->rc;
            }
            //if brother's two children are black, let brother become red, move me to the parent
            if (brother->lc->color == Black && brother->rc->color == Black)
            {
                brother->color = Red;
                me = me->parent;
            }
            else
            {
                //if brother's right is black, change the node, make left black
                if (brother->rc->color == Black)
                {
                    brother->lc->color = Black;
                    brother->color = Red;
                    right_rotate(brother);
                    brother = brother->parent;
                }
                brother->color = me->parent->color;
                me->parent->color = Black;
                brother->rc->color = Black;
                left_rotate(me->parent);
                me = root;
            }
        }
        else
        {
            //the same as the left condition
            Rnode *brother = me->parent->lc;
            if (brother->color == Red)
            {
                me->parent->color = Red;
                brother->color = Black;
                right_rotate(me->parent);
                brother = me->parent->lc;
            }
            if (brother->rc->color == Black && brother->lc->color == Black)
            {
                brother->color = Red;
                me = me->parent;
            }
            else
            {
                if (brother->lc->color == Black)
                {
                    brother->rc->color = Black;
                    brother->color = Red;
                    left_rotate(brother);
                    brother = brother->parent;
                }
                brother->color = me->parent->color;
                me->parent->color = Black;
                brother->lc->color = Black;
                right_rotate(me->parent);
                me = root;
            }
        }
    }
    me->color = Black;
}

void RBT::left_rotate(Rnode *node)
{
    //if nothing right, return
    if (node->rc == nil)
    {
        return;
    }
    Rnode *rc = node->rc;
    //调整size大小
    node->size -= rc->rc->size + 1;
    rc->size += node->lc->size + 1;

    rc->parent = node->parent;
    //if node==root, reset root
    if (node == root)
    {
        root = rc;
    }
    else
    {
        if (node == node->parent->rc)
        {
            node->parent->rc = rc;
        }
        else
        {
            node->parent->lc = rc;
        }
    }
    node->rc = rc->lc;
    rc->lc->parent = node;
    rc->lc = node;
    node->parent = rc;
}

void RBT::right_rotate(Rnode *node)
{
    //if nothing left, return
    if (node->lc == nil)
    {
        return;
    }
    Rnode *lc = node->lc;
    //调整size大小
    node->size -= lc->lc->size + 1;
    lc->size += node->rc->size + 1;

    lc->parent = node->parent;
    //if node==root, reset root
    if (node == root)
    {
        root = lc;
    }
    else
    {
        if (node == node->parent->lc)
        {
            node->parent->lc = lc;
        }
        else
        {
            node->parent->rc = lc;
        }
    }
    node->lc = lc->rc;
    lc->rc->parent = node;
    lc->rc = node;
    node->parent = lc;
}